#include <vector>
#include <unordered_set>
#include <queue>
#include "distance.hpp"

namespace GraphGradient {

    int search_enter_point(std::vector<std::vector<int>> &graph, float *query, float *base, int enter_point, int dim) {
        std::unordered_set<int> visited_element;
        std::priority_queue<distance_cmp, std::vector<distance_cmp>, distance_min_heap> candidates;
        std::priority_queue<distance_cmp, std::vector<distance_cmp>, distance_max_heap> nearest_neighbors;
        {
            visited_element.insert(enter_point);
            double tmp_dist = l2distance(query, base + enter_point * dim, dim);
            candidates.emplace(tmp_dist, enter_point);
            nearest_neighbors.emplace(tmp_dist, enter_point);
        }


        while (candidates.size() > 0) {
            distance_cmp cand = candidates.top();
            candidates.pop();
            distance_cmp far_ele = nearest_neighbors.top();
            if (cand.distance_ > far_ele.distance_) {
                break;
            }
            for (std::vector<int>::iterator iter = graph[cand.idx_].begin(); iter != graph[cand.idx_].end(); ++iter) {
                int neighbor = *iter;
                if (visited_element.find(neighbor) == visited_element.end()) { //can not find
                    visited_element.insert(neighbor);
                    far_ele = nearest_neighbors.top();
                    double neighbor_dist = l2distance(query, base + neighbor * dim, dim);
                    if (l2distance(query, base + neighbor * dim, dim) < far_ele.distance_ ||
                        (int) nearest_neighbors.size() < 1) {
                        candidates.emplace(neighbor_dist, neighbor);
                        nearest_neighbors.emplace(neighbor_dist, neighbor);
                        while ((int) nearest_neighbors.size() > 1) {
                            nearest_neighbors.pop();
                        }
                    }
                }
            }

        }
        while (nearest_neighbors.size() > 10) {
            nearest_neighbors.pop();
        }

        std::vector<int> res;
        while (nearest_neighbors.size() > 0) {
            distance_cmp tmp = nearest_neighbors.top();
            nearest_neighbors.pop();
            res.push_back(tmp.idx_);
        }
        int res_size = res.size();
        return res[res_size - 1];
    }

    float *get_centroid(float *base, int n_base, int vec_dim) {
        double *tmp_centroid = new double[vec_dim];
        float *centroid = new float[vec_dim];
        float *tmp_item = base;
        for (int i = 0; i < vec_dim; i++) {
            tmp_centroid[i] = 0;
        }
        for (int i = 0; i < n_base; i++) {
            for (int j = 0; j < vec_dim; j++) {
                tmp_centroid[j] += tmp_item[j];
            }
            tmp_item += vec_dim;
        }
        for (int i = 0; i < vec_dim; i++) {
            centroid[i] = (float) (tmp_centroid[i] / n_base);
        }
        delete[] tmp_centroid;
        return centroid;
    }

    int get_ep_centroid(std::vector<std::vector<int>> &graph, float *base, int random_ep, int dim) {
        int n_base = graph.size();
        float *centroid = get_centroid(base, n_base, dim);

        int res = search_enter_point(graph, centroid, base, random_ep, dim);
        return res;
    }
}
